from typing import Dict, List, Optional
import openai
from .config import AIConfig

class AIModel:
    """AI模型管理类，负责与OpenAI API通信和模型调用"""
    def __init__(self, config: Optional[AIConfig] = None):
        self.config = config if config else AIConfig()
        self.setup_api()
    
    def setup_api(self) -> None:
        """设置OpenAI API配置"""
        api_key = self.config.api_key
        if not api_key:
            raise ValueError("API密钥不能为空")
        openai.api_key = api_key
    
    def generate_response(self, prompt: str, **kwargs) -> Optional[str]:
        """生成AI响应"""
        try:
            # 获取模型配置
            model = kwargs.get('model', self.config.model_name)
            temperature = float(kwargs.get('temperature', self.config.get_config('temperature')))
            max_tokens = int(kwargs.get('max_tokens', self.config.get_config('max_tokens')))
            
            # 调用API
            response = openai.ChatCompletion.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                max_tokens=max_tokens
            )
            
            # 提取响应文本
            if response and hasattr(response, 'choices') and response.choices:
                if isinstance(response.choices[0].message, dict):
                    return response.choices[0].message["content"]
                return response.choices[0].message.content
            return None
            
        except Exception as e:
            print(f'生成响应失败: {e}')
            raise
    
    def validate_api_key(self, api_key: str) -> bool:
        """验证API密钥是否有效"""
        try:
            openai.api_key = api_key
            models = openai.Model.list()
            return bool(models)
        except:
            return False
    
    def get_available_models(self) -> List[str]:
        """获取可用的模型列表"""
        try:
            models = openai.Model.list()
            return [model.id for model in models.data if 'gpt' in model.id.lower()]
        except:
            return [self.config.default_config['model']]
    
    def update_model_config(self, config: Dict) -> bool:
        """更新模型配置"""
        try:
            for key, value in config.items():
                if not self.config.set_config(key, value):
                    return False
            self.setup_api()
            return True
        except:
            return False